import os
from loguru import logger
import torch
from torch.utils.data import DataLoader
from torchvision.datasets import MNIST
from torchvision.transforms import transforms

from model.LeNet import LeNet5
from config import cfg


def inference(config, data_iter):
    model = LeNet5()
    model.to(cfg.device)
    if os.path.exists(config.model_save_path):
        model_params = torch.load(config.model_save_path)["model_state_dict"]
        model.load_state_dict(model_params)
        model.eval()

    test_loader = DataLoader(data_iter, batch_size=config.batch_size, shuffle=False)
    with torch.no_grad():
        acc_sum, n = 0.0, 0
        for batch_idx, (data, target) in enumerate(test_loader):
            data, target = data.to(config.device), target.to(config.device)
            logits = model(data)
            acc_sum += (logits.argmax(1) == target).float().sum().item()
            n += len(target)
            logger.info(f"Accuracy in test data is : {acc_sum / n}")


if __name__ == "__main__":
    mnist_test = MNIST(root=cfg.dataset_dir,
                       train=False,
                       download=True,
                       transform=transforms.ToTensor())
    inference(config=cfg, data_iter=mnist_test)
